﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using System.Windows.Forms;
using Ast.Core.Forms;
using Ast.Core.Services;
using Ast.Core.Services.Extensions;
using Ast.Core.Services.ObjectDefinitions;
using Ast.Framework.Forms;
using Castle.Core.Logging;

namespace Ast
{
    [FormMetadata(FormType.Process, IsTabForm = true)]
    public partial class ProcessForm : Form
    {
        private readonly IPspCidTableService _pspCidTableService;
        private readonly IObjectTypeService _objectTypeService;

        public ProcessForm(IPspCidTableService pspCidTableService, IObjectTypeService objectTypeService)
        {
            _pspCidTableService = pspCidTableService;
            _objectTypeService = objectTypeService;

            InitializeComponent();

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        private void ProcessForm_Load(object sender, EventArgs e)
        {
            lvProcesses.ApplyColumnSorter();
            btnReload.PerformClick();
        }

        private async void btnReload_Click(object sender, EventArgs e)
        {
            var items = new List<ListViewItem>();

            await Task.Run(() => _pspCidTableService.ListProcesses((address, process, path) =>
            {
                var processViewItem = new ProcessViewItem
                {
                    Pid = process.UniqueProcessId,
                    Name = process.ImageFileName,
                    Path = path,
                    EProcessAddress = address,
                    EProcess = process
                };

                var listViewItem = new ListViewItem(new[]
                {
                    processViewItem.Pid.ToString(),
                    processViewItem.Name,
                    processViewItem.Path,
                    string.Format("0x{0:x}", processViewItem.EProcessAddress.ToInt32())
                })
                {
                    Tag = processViewItem
                };

                items.Add(listViewItem);

                return true; // continue search
            })).ConfigureAwait(false);

            Logger.InfoFormat("found {0} processes", items.Count);

            // update view
            Invoke(new Action(() =>
            {
                lvProcesses.BeginUpdate();
                lvProcesses.Items.Clear();
                lvProcesses.Items.AddRange(items.ToArray());
                lvProcesses.AutoResizeColumns();
                lvProcesses.EndUpdate();
            }));
        }

        private void contextMenuStrip1_Opening(object sender, System.ComponentModel.CancelEventArgs e)
        {
            e.Cancel = lvProcesses.SelectedItems.Count == 0;
        }

        private async void viewThreadsToolStripMenuItem_Click(object sender, EventArgs e)
        {
            if (lvProcesses.SelectedItems.Count == 0) return;
            var processViewItem = (ProcessViewItem) lvProcesses.SelectedItems[0].Tag;

            // seek threads
            var items = new List<ThreadViewItem>();
            await Task.Run(() => _pspCidTableService.ListThreads((address, thread) =>
            {
                if (thread.ClientId.UniqueProcess == processViewItem.Pid)
                {
                    items.Add(new ThreadViewItem
                    {
                        Pid = thread.ClientId.UniqueProcess,
                        Tid = thread.ClientId.UniqueThread,
                        EThreadAddress = address,
                        Teb = thread.Teb,
                        StartAddress = thread.StartAddress,
                        Win32StartAddress = thread.Win32StartAddress,
                        EThread = thread
                    });
                }
                return true; // continue search
            })).ConfigureAwait(false);

            Logger.InfoFormat("found {0} threads for pid {1}", items.Count, processViewItem.Pid);
            if (items.Count == 0) return;

            // show threads
            var threadForm = new ListViewForm<ThreadViewItem>(new ListViewFormDataSource<ThreadViewItem>
            {
                Title = string.Format("Threads for {0}({1})", processViewItem.Name, items.Count),
                Columns = new List<ListViewColumn>
                {
                    new ListViewColumn {Name = "Pid", Tag = "number"},
                    new ListViewColumn {Name = "Tid", Tag = "number"},
                    new ListViewColumn {Name = "Teb", Tag = null},
                    new ListViewColumn {Name = "EThread", Tag = null},
                    new ListViewColumn {Name = "Start Address", Tag = null},
                    new ListViewColumn {Name = "Win32 Start Address", Tag = null},
                },
                Data = items,
                DataGeters = new List<Func<ThreadViewItem, ListViewColumn, string>>
                {
                    (item, column) => item.Pid.ToString(),
                    (item, column) => item.Tid.ToString(),
                    (item, column) => string.Format("0x{0:x8}", item.Teb.ToInt32()),
                    (item, column) => string.Format("0x{0:x8}", item.EThreadAddress.ToInt32()),
                    (item, column) => string.Format("0x{0:x8}", item.StartAddress.ToInt32()),
                    (item, column) => string.Format("0x{0:x8}", item.Win32StartAddress.ToInt32()),
                }
            });
            BeginInvoke(new Action(() => threadForm.ShowDialog()));
        }

        private async void viewHandlesToolStripMenuItem_Click(object sender, EventArgs e)
        {
            if (lvProcesses.SelectedItems.Count == 0) return;
            var processViewItem = (ProcessViewItem) lvProcesses.SelectedItems[0].Tag;

            // seek handles
            var items = new List<HandleViewItem>();
            await Task.Run(() =>
            {
                var objectTypes = _objectTypeService.ListObjectTypesInCache(false);

                var offset = Marshal.SizeOf<ObjectHeader>(); // object offset is 0x18 since all entries in the table are pointing to _OBJECT_HEADER
                _pspCidTableService.ListObjects(processViewItem.EProcess.ObjectTable, (objPtr, objType) =>
                {
                    var type = objectTypes.FirstOrDefault(x => x.ObjectType.Address == objType);
                    if (type != null)
                    {
                        var item = new HandleViewItem
                        {
                            Pid = processViewItem.Pid,
                            ObjectAddress = objPtr,
                            Type = type,
                            PointerCount = _objectTypeService.GetObjectHeader(objPtr).Object.PointerCount,
                            Name = _objectTypeService.GetObjectNameByType(objPtr, type),
                        };

                        items.Add(item);
                    }
                    else
                    {
                        Logger.WarnFormat("not found object type for object 0x{0:x}", objPtr.ToInt32());
                    }

                    return true;
                }, offset);
            }).ConfigureAwait(false);

            Logger.InfoFormat("found {0} handles for pid {1}", items.Count, processViewItem.Pid);
            if (items.Count == 0) return;

            // show handles
            var handleForm = new ListViewForm<HandleViewItem>(new ListViewFormDataSource<HandleViewItem>
            {
                Title = string.Format("Handles for {0}({1})", processViewItem.Name, items.Count),
                Columns = new List<ListViewColumn>
                {
                    new ListViewColumn {Name = "Pid", Tag = "number"},
                    new ListViewColumn {Name = "Object", Tag = null},
                    new ListViewColumn {Name = "Type", Tag = null},
                    new ListViewColumn {Name = "Type Index", Tag = "number"},
                    new ListViewColumn {Name = "Pointer Count", Tag = "number"},
                    new ListViewColumn {Name = "Name", Tag = null},
                },
                Data = items,
                DataGeters = new List<Func<HandleViewItem, ListViewColumn, string>>
                {
                    (item, column) => item.Pid.ToString(),
                    (item, column) => string.Format("0x{0:x}", item.ObjectAddress.ToInt32()),
                    (item, column) => item.Type.Name,
                    (item, column) => item.Type.ObjectType.Object.Index.ToString(),
                    (item, column) => item.PointerCount.ToString(),
                    (item, column) => item.Name,
                }
            });
            BeginInvoke(new Action(() => handleForm.ShowDialog()));
        }
    }

    public class ProcessViewItem
    {
        public int Pid { get; set; }
        public string Name { get; set; }
        public string Path { get; set; }
        public IntPtr EProcessAddress { get; set; }

        public EProcess EProcess { get; set; }
    }

    public class ThreadViewItem
    {
        public int Pid { get; set; }
        public int Tid { get; set; }
        public IntPtr EThreadAddress { get; set; }
        public IntPtr Teb { get; set; }
        public IntPtr StartAddress { get; set; }
        public IntPtr Win32StartAddress { get; set; }

        public EThread EThread { get; set; }
    }

    public class HandleViewItem
    {
        public int Pid { get; set; }
        public IntPtr ObjectAddress { get; set; }
        public ObjectTypeData Type { get; set; }
        public string Name { get; set; }
        public int PointerCount { get; set; }
    }
}